#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File    :   1_read_dureader_data.py
@Contact :   cr_lgdx@163.com
@Author  :   lduml
@Modify Time      @Version    @Desciption
------------      --------    -----------
2020/3/31 14:44    1.0         读取文件，修改格式，不进行shuffle
"""
import json
import pandas as pd
import random
import argparse
import os

# 设置固定的随机种子
random.seed(1234)

parser = argparse.ArgumentParser()
parser.add_argument("--train_file", help="train_file",type=str, default='train.json')
parser.add_argument('--dev_file', help="dev_file",type=str, default='dev.json')
parser.add_argument('--test_file', help="dev_file",type=str, default='test.json')
parser.add_argument('--out_put_dir', help="dev_file",type=str, default='out_put_dir')



def load_train_dev_file(file_dir, style):
    # 训练数据
    with open(file_dir, encoding='utf-8') as f:
        train_data = json.load(f)
    paragraphs = train_data['data'][0]['paragraphs']
    for item in paragraphs:
        try:
            item['question'] = item['qas'][0]['question']
            item['id'] = item['qas'][0]['id']
            item['answers_text'] = item['qas'][0]['answers'][0]['text']
            item['answer_start'] = item['qas'][0]['answers'][0]['answer_start']
            del item['qas']
        except Exception as e:
            print(e)
    # 保存数据
    train_data_pd = pd.DataFrame.from_records(paragraphs)
    train_data_pd.to_csv(os.path.join(args.out_put_dir, style + '.csv'), index=False, encoding='utf-8')
    print(os.path.join(args.out_put_dir, style + '.csv'))


def load_test_data(file_dir):
    # 测试集数据处理
    with open(file_dir, encoding='utf-8') as f:
        test_data = json.load(f)
    paragraphs = test_data['data'][0]['paragraphs']
    for item in paragraphs:
        try:
            item['question'] = item['qas'][0]['question']
            item['id'] = item['qas'][0]['id']
            del item['qas']
        except Exception as e:
            print(e)
            print(item)
    test_data_pd = pd.DataFrame.from_records(paragraphs)
    test_data_pd.to_csv(os.path.join(args.out_put_dir, 'test_data.csv'), index=False, encoding='utf-8')
    print(os.path.join(args.out_put_dir, 'test_data.csv'))


if __name__=='__main__':
    args = parser.parse_args()
    train_file = args.train_file
    dev_file = args.dev_file
    test_file = args.test_file

    load_train_dev_file(train_file, 'train_data')
    load_train_dev_file(dev_file, 'dev_data')
    load_test_data(test_file)
